"""
3D Navier-Stokes Solver

Implements the 3D vector stream function-vorticity formulation for incompressible flow
on L-shaped domains. Solves the coupled system:

1. Poisson equations: ∇²ψₖ = -ωₖ for k ∈ {x,y,z} (vector stream function)
2. Vorticity transport: ∂ω/∂t + u·∇ω - ω·∇u = (1/Re)∇²ω (includes stretching term)

where u = ∇ × ψ are the derived velocities from the vector stream function.

The 3D formulation extends the 2D case to handle:
- Vector stream functions: ψ = (ψₓ, ψᵧ, ψᵤ)
- Vector vorticity: ω = (ωₓ, ωᵧ, ωᵤ)
- Vorticity stretching terms: (ω·∇)u in the transport equation
- 3D boundary conditions including moving lid in z-direction
"""

"""
    solve_lid_driven_3d_impl(options::LidDrivenOptions) -> LidDrivenResult

Core implementation of the 3D lid-driven cavity solver using vector stream function-vorticity formulation.

# Algorithm Overview
1. Initialize 3D domain, grid, and vector fields
2. Create Poisson solvers for three stream function components
3. Time-stepping loop:
   a. Apply boundary conditions to vector stream function
   b. Solve three Poisson equations: ∇²ψₖ = -ωₖ
   c. Compute velocities from vector stream function: u = ∇ × ψ
   d. Update vorticity transport with stretching: ∂ω/∂t + u·∇ω - ω·∇u = (1/Re)∇²ω
   e. Apply vorticity boundary conditions
   f. Apply under-relaxation for stability
   g. Check convergence
4. Return solution and diagnostics

# Mathematical Formulation
The vector stream function ψ = (ψₓ, ψᵧ, ψᵤ) and velocity u = ∇ × ψ:
- u = ∂ψᵤ/∂y - ∂ψᵧ/∂z
- v = ∂ψₓ/∂z - ∂ψᵤ/∂x
- w = ∂ψᵧ/∂x - ∂ψₓ/∂y

Vorticity ω = ∇ × u satisfies:
∂ω/∂t + u·∇ω - ω·∇u = (1/Re)∇²ω

Boundary conditions:
- Moving lid (y=1, x≤0): u=U_lid, v=w=0, ψᵤ=x+1, ψₓ=ψᵧ=0
- No-slip walls: u=v=w=0, ψₓ=ψᵧ=ψᵤ=0

# Arguments
- `options::LidDrivenOptions`: Solver configuration

# Returns
- `LidDrivenResult`: Complete 3D solution with vector fields and diagnostics
"""
function solve_lid_driven_3d_impl(options::LidDrivenOptions)
    
    if options.verbose
        println("🌪️  Starting 3D Lid-Driven Cavity Solver (Vector Stream Function)")
        println("Configuration: N=$(options.n), Re=$(options.Re), solver=$(options.solver)")
        println("WARNING: 3D solver is computationally intensive!")
    end
    
    start_time = time()
    
    # ================================================================
    # 1. Initialize 3D domain and grid  
    # ================================================================
    
    x, y, z, h = build_grid(options.n, Val(3))
    domain = create_lshape_domain(options.n, Val(3))
    nx, ny, nz = length(x), length(y), length(z)
    
    if options.verbose
        print_domain_info(domain)
        print_grid_info(x, y, z, options.Re)
    end
    
    # ================================================================
    # 2. Initialize vector fields
    # ================================================================
    
    # Vector stream function: ψ = (ψₓ, ψᵧ, ψᵤ)
    ψ = zeros(Float64, nx, ny, nz, 3)  # 4D array: [x, y, z, component]
    
    # Vector vorticity: ω = (ωₓ, ωᵧ, ωᵤ)
    ω = zeros(Float64, nx, ny, nz, 3)
    
    # Velocity components
    u = zeros(Float64, nx, ny, nz)
    v = zeros(Float64, nx, ny, nz)
    w = zeros(Float64, nx, ny, nz)
    
    # Set NaN in solid regions for visualization
    for comp in 1:3
        ψ[:, :, :, comp][.!domain.fluid_mask] .= NaN
        ω[:, :, :, comp][.!domain.fluid_mask] .= NaN
    end
    u[.!domain.fluid_mask] .= 0.0
    v[.!domain.fluid_mask] .= 0.0
    w[.!domain.fluid_mask] .= 0.0
    
    # ================================================================
    # 3. Create Poisson solvers for each stream function component
    # ================================================================
    
    # We need three separate Poisson solvers for ∇²ψₖ = -ωₖ
    poisson_solvers = []
    for k in 1:3
        solver = create_poisson_solver(domain, h, options)
        push!(poisson_solvers, solver)
    end
    
    if options.verbose
        @info "Created $(length(poisson_solvers)) Poisson solvers for vector components"
    end
    
    # ================================================================
    # 4. Auto-adjust relaxation parameters for 3D stability
    # ================================================================
    
    α_ψ, α_ω = compute_auto_relaxation_3d(options.α_ψ, options.α_ω, options.n, options.solver, options.verbose)
    
    # ================================================================
    # 5. Time step and stability setup
    # ================================================================
    
    dt = compute_stable_timestep_3d(h, options.Re, options.dt, options.verbose)
    
    # ================================================================
    # 6. Initialize diagnostics
    # ================================================================
    
    convergence_history = Float64[]
    timing_breakdown = (poisson=0.0, velocity=0.0, vorticity=0.0, boundary=0.0)
    
    # ================================================================
    # 7. Main time-stepping loop
    # ================================================================
    
    converged = false
    final_residual = Inf
    
    for step in 1:options.max_steps
        step_start = time()
        
        # ==============================================================
        # 7a. Apply boundary conditions to vector stream function
        # ==============================================================
        
        t_bc = time()
        apply_boundary_conditions!(ψ, ω, u, v, w, domain, x, y, z, options.U_lid, h)
        timing_breakdown = merge(timing_breakdown, (boundary = timing_breakdown.boundary + time() - t_bc,))
        
        # ==============================================================
        # 7b. Solve three Poisson equations: ∇²ψₖ = -ωₖ 
        # ==============================================================
        
        t_poisson = time()
        
        # Store old values for under-relaxation
        ψ_old = copy(ψ)
        
        # Solve for each stream function component
        for k in 1:3
            # Create RHS field: -ωₖ
            rhs_field = -ω[:, :, :, k]
            
            # Solve ∇²ψₖ = rhs_field
            psi_component = view(ψ, :, :, :, k)
            solve_result = solve_poisson!(poisson_solvers[k], rhs_field, psi_component)
            
            if !solve_result.converged && options.verbose
                @warn "Poisson solver for component $k did not converge (iter=$(solve_result.iterations))"
            end
        end
        
        # Apply under-relaxation to stream function
        for k in 1:3
            ψ[:, :, :, k] .= ψ_old[:, :, :, k] .+ α_ψ .* (ψ[:, :, :, k] .- ψ_old[:, :, :, k])
        end
        
        timing_breakdown = merge(timing_breakdown, (poisson = timing_breakdown.poisson + time() - t_poisson,))
        
        # ==============================================================
        # 7c. Compute velocities from vector stream function: u = ∇ × ψ
        # ==============================================================
        
        t_velocity = time()
        
        # Extract stream function components for clarity
        ψx = view(ψ, :, :, :, 1)
        ψy = view(ψ, :, :, :, 2) 
        ψz = view(ψ, :, :, :, 3)
        
        # Compute velocities using curl operation
        compute_velocities_from_stream_3d!(u, v, w, ψx, ψy, ψz, domain, h)
        
        # Adaptive time step based on maximum velocity
        u_max = maximum(abs.(u[domain.fluid_mask]))
        v_max = maximum(abs.(v[domain.fluid_mask]))
        w_max = maximum(abs.(w[domain.fluid_mask]))
        vel_max = max(u_max, v_max, w_max, 1e-12)
        
        # Update time step based on CFL condition
        cfl_factor = 0.25
        dt_cfl = cfl_factor * h / vel_max
        dt_diff = 0.10 * h^2 * options.Re  # Diffusion stability
        dt = min(dt_cfl, dt_diff)
        
        timing_breakdown = merge(timing_breakdown, (velocity = timing_breakdown.velocity + time() - t_velocity,))
        
        # ==============================================================
        # 7d. Update vorticity transport with stretching terms
        # ==============================================================
        
        t_vorticity = time()
        
        ω_old = copy(ω)
        
        # Update each vorticity component with transport equation
        update_vorticity_transport_3d!(ω, u, v, w, domain, h, dt, options.Re)
        
        # Apply vorticity boundary conditions
        apply_boundary_conditions!(ψ, ω, u, v, w, domain, x, y, z, options.U_lid, h)
        
        # Apply under-relaxation to vorticity
        for k in 1:3
            ω[:, :, :, k] .= ω_old[:, :, :, k] .+ α_ω .* (ω[:, :, :, k] .- ω_old[:, :, :, k])
        end
        
        timing_breakdown = merge(timing_breakdown, (vorticity = timing_breakdown.vorticity + time() - t_vorticity,))
        
        # ==============================================================
        # 7e. Convergence check
        # ==============================================================
        
        # Compute residual based on vorticity change
        residual = 0.0
        for k in 1:3
            ω_change = ω[:, :, :, k] .- ω_old[:, :, :, k]
            residual += maximum(abs.(ω_change[domain.fluid_mask]))
        end
        residual /= (options.n^1.5)  # 3D normalization
        
        push!(convergence_history, residual)
        
        # Progress reporting
        if options.verbose && (step % 10 == 0 || step <= 3)
            ω_norms = [norm(ω[:, :, :, k][domain.fluid_mask]) for k in 1:3]
            @printf("Step %4d: residual=%.3e, dt=%.2e, |u|max=%.3e, |ω|=(%.2e,%.2e,%.2e), time=%.1fs\n",
                   step, residual, dt, vel_max, ω_norms..., time()-step_start)
        end
        
        # Check convergence
        if residual < options.tol
            converged = true
            final_residual = residual
            if options.verbose
                @info "✓ Converged after $step steps: residual = $residual"
            end
            break
        end
        
        final_residual = residual
    end
    
    total_time = time() - start_time
    
    if !converged
        @warn "⚠️  Did not converge within $(options.max_steps) steps (final residual: $final_residual)"
    end
    
    if options.verbose
        @info "Total solve time: $(round(total_time, digits=2))s"
        @info "Time breakdown: Poisson=$(round(timing_breakdown.poisson, digits=1))s, " *
              "Velocity=$(round(timing_breakdown.velocity, digits=1))s, " *
              "Vorticity=$(round(timing_breakdown.vorticity, digits=1))s, " *
              "BC=$(round(timing_breakdown.boundary, digits=1))s"
    end
    
    # ================================================================
    # 8. Return results
    # ================================================================
    
    return LidDrivenResult(
        ψ = ψ,
        ω = ω,
        u = u,
        v = v,
        w = w,
        x = x,
        y = y,
        z = z,
        converged = converged,
        iterations = length(convergence_history),
        residual_norm = final_residual,
        solve_time = total_time,
        convergence_history = convergence_history,
        options = options
    )
end

"""
    update_vorticity_transport_3d!(ω, u, v, w, domain, h, dt, Re)

Update vorticity using the 3D transport equation with stretching terms.

The 3D vorticity transport equation is:
∂ω/∂t + u·∇ω - ω·∇u = (1/Re)∇²ω

where:
- u·∇ω: convection term (vorticity advection)
- ω·∇u: stretching term (vorticity stretching/tilting)
- (1/Re)∇²ω: diffusion term

# Arguments
- `ω::Array{Float64,4}`: Vorticity field (modified in-place)
- `u,v,w::Array{Float64,3}`: Velocity components
- `domain::LShapeDomain{3}`: Domain geometry
- `h::Float64`: Grid spacing
- `dt::Float64`: Time step
- `Re::Float64`: Reynolds number
"""
function update_vorticity_transport_3d!(ω::Array{Float64,4}, 
                                      u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                                      domain::LShapeDomain{3}, h::Float64, dt::Float64, Re::Float64)
    nx, ny, nz = size(u)
    
    # Create temporary arrays for the new vorticity values
    ω_new = copy(ω)
    
    # Update interior nodes only
    @inbounds for k in 2:nz-1, j in 2:ny-1, i in 2:nx-1
        if !domain.interior_mask[i, j, k]
            continue  # Skip non-interior nodes
        end
        
        # Update each vorticity component
        for comp in 1:3
            # Current vorticity value
            ω_curr = ω[i, j, k, comp]
            
            # Convection term: u·∇ω
            convection = upwind_convection_3d(ω[:, :, :, comp], u, v, w, i, j, k, h)
            
            # Diffusion term: ∇²ω
            diffusion = laplacian_3d(ω[:, :, :, comp], i, j, k, h)
            
            # Stretching term: (ω·∇)u component
            stretch_x, stretch_y, stretch_z = vorticity_stretching_3d(ω[:, :, :, 1], ω[:, :, :, 2], ω[:, :, :, 3],
                                                                     u, v, w, i, j, k, h)
            
            # Select appropriate stretching component
            if comp == 1
                stretching = stretch_x
            elseif comp == 2  
                stretching = stretch_y
            else
                stretching = stretch_z
            end
            
            # Explicit Euler time stepping: 
            # ∂ω/∂t = -u·∇ω + (ω·∇)u + (1/Re)∇²ω
            rhs = -convection + stretching + diffusion / Re
            ω_new[i, j, k, comp] = ω_curr + dt * rhs
        end
    end
    
    # Copy updated values back to original array
    ω .= ω_new
end

"""
    compute_stable_timestep_3d(h::Float64, Re::Float64, dt_user::Union{Float64, Nothing}, verbose::Bool) -> Float64

Compute stable time step for 3D explicit time integration.

# Arguments
- `h::Float64`: Grid spacing
- `Re::Float64`: Reynolds number  
- `dt_user::Union{Float64, Nothing}`: User-specified time step (if any)
- `verbose::Bool`: Whether to print diagnostic information

# Returns
- `dt::Float64`: Stable time step
"""
function compute_stable_timestep_3d(h::Float64, Re::Float64, dt_user::Union{Float64, Nothing}, verbose::Bool)
    # Conservative stability factors for 3D
    cfl_factor = 0.25      # CFL condition: dt ≤ C·h/|u|_max
    diffusion_factor = 0.10 # Diffusion condition: dt ≤ C·h²·Re (3D needs smaller factor)
    
    # Stability limits
    dt_cfl = cfl_factor * h           # CFL limit (assuming |u|_max ~ 1)
    dt_diffusion = diffusion_factor * h^2 * Re  # Diffusion limit
    
    # Most restrictive condition
    dt_stable = min(dt_cfl, dt_diffusion)
    
    # Use user value if specified and not too large
    dt = if dt_user !== nothing
        if dt_user <= dt_stable
            dt_user
        else
            if verbose
                @warn "User dt=$dt_user exceeds stability limit $dt_stable, using stable value"
            end
            dt_stable
        end
    else
        dt_stable
    end
    
    if verbose
        @info "Time step selection: dt=$dt (CFL limit=$dt_cfl, diffusion limit=$dt_diffusion)"
        if dt_diffusion < dt_cfl
            @info "Time step is diffusion-limited (small Re or fine grid)"
        else
            @info "Time step is CFL-limited (large velocities expected)"
        end
    end
    
    return dt
end

"""
    compute_auto_relaxation_3d(α_ψ_user::Union{Float64, Nothing}, α_ω_user::Union{Float64, Nothing},
                              n::Int, solver::Symbol, verbose::Bool) -> Tuple{Float64, Float64}

Auto-compute relaxation parameters for 3D stability.

# Arguments
- `α_ψ_user, α_ω_user::Union{Float64, Nothing}`: User-specified relaxation parameters
- `n::Int`: Grid size parameter
- `solver::Symbol`: Linear solver type
- `verbose::Bool`: Whether to print information

# Returns
- `Tuple{Float64, Float64}`: (α_ψ, α_ω) relaxation parameters
"""
function compute_auto_relaxation_3d(α_ψ_user::Union{Float64, Nothing}, α_ω_user::Union{Float64, Nothing},
                                   n::Int, solver::Symbol, verbose::Bool)
    
    # Conservative defaults for 3D (more restrictive than 2D)
    α_ψ_default = 0.15  # Stream function relaxation  
    α_ω_default = 0.45  # Vorticity relaxation
    
    # Adjust based on grid size (finer grids need more relaxation)
    if n >= 32
        α_ψ_default *= 0.8
        α_ω_default *= 0.9
    end
    
    # Adjust based on solver type
    if solver == :julia
        # Built-in solvers may need more relaxation
        α_ψ_default *= 0.9
    end
    
    # Use user values if provided
    α_ψ = something(α_ψ_user, α_ψ_default)
    α_ω = something(α_ω_user, α_ω_default)
    
    # Clamp to reasonable ranges
    α_ψ = clamp(α_ψ, 0.05, 0.5)
    α_ω = clamp(α_ω, 0.1, 0.8)
    
    if verbose
        if α_ψ_user === nothing || α_ω_user === nothing
            @info "Auto-selected relaxation: α_ψ=$α_ψ, α_ω=$α_ω (based on n=$n, solver=$solver)"
        else
            @info "Using specified relaxation: α_ψ=$α_ψ, α_ω=$α_ω"
        end
    end
    
    return α_ψ, α_ω
end